import argparse
import time
from urllib import parse
import requests
import pickle
import os
from PIL import Image
from requests import Response
from rich.progress import Progress, SpinnerColumn, BarColumn, DownloadColumn, TaskProgressColumn, TimeElapsedColumn, \
    TimeRemainingColumn, TextColumn, TransferSpeedColumn, TaskID
from console import *
from qrcode import QRCode, constants
from qrcode.image.styledpil import StyledPilImage
from qrcode.image.styles.moduledrawers import RoundedModuleDrawer
import re
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from threading import Event
from typing import Iterable
from urllib.request import urlopen, Request


class BiliBiLi(object):

    """
    请求对象
    """

    # 工具标题
    name = "BiliBili视频下载器"

    # Cookies 保存路径
    cookies_save_path = os.path.join(os.getcwd(), "cookies", "bilibili.cookies")

    # 请求Session
    session = None

    # 登录二维码保存位置
    login_qr_code_path = os.path.join(os.getcwd(), "code", "login.png")

    # 公共URL
    base_url = "https://"

    # 视频url
    url = None

    # 请求参数
    data = {}

    # 请求头
    headers = {
        "referer": "https://www.bilibili.com",
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36"
    }

    # Cookies
    cookies = {}

    # 请求超时
    timeout = 5

    # 是否已经登陆
    is_login = False

    # 是否大会员
    is_vip = False

    # 视频bvid
    bvid = None

    # 视频下载路径
    save_path = os.path.join(os.getcwd(), "download")

    # 视频标题
    title = None

    # 视频清晰度列表
    definition_data = {
        6: "240P 极速",
        16: "360P 流畅",
        32: "480P 清晰",
        64: "720P 高清",
        74: "720P60 高清",
        80: "1080P 高清",
        112: "1080P+ 高清",
        116: "1080P60 高清",
        120: "4K 超清"
    }

    # 最大下载任务
    max_works = 5

    # 完成是否关机
    is_exit = False

    # 事件处理
    done_event = Event()

    # 命令行参数
    args = None

    # 进度条
    progress = Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]{task.fields[filename]}", justify="left"),
        BarColumn(bar_width=None),
        "[progress.percentage]{task.percentage:>3.1f}%",
        "•",
        DownloadColumn(),
        "•",
        TransferSpeedColumn(),
        "•",
        TimeRemainingColumn(),
    )

    def __init__(self):
        """
        请求对象
        """
        banner = r"""
                    \\         //
                     \\       //
                #####################     ________   ___   ___        ___   ________   ___   ___        ___
                ##                 ##    |\   __  \ |\  \ |\  \      |\  \ |\   __  \ |\  \ |\  \      |\  \
                ##    //     \\    ##    \ \  \|\ /_\ \  \\ \  \     \ \  \\ \  \|\ /_\ \  \\ \  \     \ \  \
                ##   //       \\   ##     \ \   __  \\ \  \\ \  \     \ \  \\ \   __  \\ \  \\ \  \     \ \  \
                ##                 ##      \ \  \|\  \\ \  \\ \  \____ \ \  \\ \  \|\  \\ \  \\ \  \____ \ \  \
                ##       www       ##       \ \_______\\ \__\\ \_______\\ \__\\ \_______\\ \__\\ \_______\\ \__\
                ##                 ##        \|_______| \|__| \|_______| \|__| \|_______| \|__| \|_______| \|__|
                #####################
                    \/         \/                               哔哩哔哩 (゜-゜)つロ 干杯~
                                                         温馨提示: 运行前请先 -h | --help 查看帮助
                """
        primary(banner)

        # 命令行获取参数
        self.parse_cmd()

        # Session
        self.session = requests.session()

        # 尝试读取本地缓存Cookies
        self.read_cookies()

        # 视频链接
        self.url = self.args.url if self.args.url else input("请输入视频链接: ")

        # 同时下载数
        self.max_works = self.args.max if self.args.max else int(input("请输入同时下载数(必须是正整数): "))

        # 是否关机
        self.is_exit = "是" if self.args.exit == "是" else "否" if self.args.exit == "否" else "是" if input("下载完成是否关机(默认: 否): ") == "是" else "否"

        # 是否登录
        if ("是" if self.args.login == "是" else "否" if self.args.login == "否" else "是" if input("是否登录以获取更好的清晰度(默认: 是): ") == "是" else "否") == "是":
            self.login()

        # 开始下载
        self.start_download()

        # 是否关机
        self.exit()

    def get(self, url: str = "", headers: dict = {}, cookies: dict = {}, timeout=5) -> Response:
        """
        GET请求
        :param url: 请求URL
        :param headers: 请求头
        :param cookies: Cookies
        :param timeout:请求超时
        :return:
        """
        if not url:
            error("请求URL不能为空")
            return
        url = self.base_url + url
        headers = {**self.headers, **headers}
        cookies = {**self.cookies, **cookies}
        self.session.cookies.update(cookies)
        timeout = timeout if timeout else self.timeout
        return self.session.get(url=url, headers=headers, cookies=cookies, timeout=timeout)

    def post(self, url: str = "", data: dict = {}, headers: dict = {}, cookies: dict = {}, timeout=5) -> Response:
        """
        POST请求
        :param url: 请求URL
        :param data: 请求参数
        :param headers: 请求头
        :param cookies: Cookies
        :param timeout:请求超时
        :return:
        """
        if not url:
            error("请求URL不能为空")
            return
        url = self.base_url + url
        data = {**self.data, **data}
        headers = {**self.headers, **headers}
        cookies = {**self.cookies, **cookies}
        self.session.cookies.update(cookies)
        timeout = timeout if timeout else self.timeout
        return self.session.post(url=url, data=data, headers=headers, cookies=cookies, timeout=timeout)

    def login(self) -> None:
        """
        登录BiliBili
        :return:
        """
        if self.is_login:
            success("登陆成功")
            return
        data = self.get("passport.bilibili.com/x/passport-login/web/qrcode/generate?source=main-fe-header").json()
        if data['code'] != 0:
            error("获取登录链接失败")
            return
        result = data['data']
        # 生成登录二维码
        self.generate_qr_code(result['url'])

        # 每2秒轮询登录链接
        for i in range(0, 160, 2):
            verify_data = self.get(f"passport.bilibili.com/x/passport-login/web/qrcode/poll?qrcode_key={result['qrcode_key']}",
                                    ).json()
            verify_data = verify_data["data"]
# https://passport.bilibili.com/x/passport-login/web/qrcode/poll?qrcode_key=0fca55d13f5edbdc57751d4a8bedaa49&source=main-fe-header
            print(verify_data)
            try:
                if verify_data['code'] == 0:
                    os.remove(self.login_qr_code_path)
                    # 从URL解析出Cookies
                    self.cookies = self.parser_url_params(verify_data['url'])
                    # 更新Session的Cookies
                    self.session.cookies.update(self.cookies)
                    self.is_login = self.verify_cookies()
                    # 保存Cookies到本地
                    self.save_cookies()
                    success("登录成功")
                    break
            except:
                data = verify_data
                if data == 86101:
                    warn("请在规定时间内进行扫码")
                elif data == 86090:
                    primary("二维码已扫码未确认")
                elif data == -1:
                    os.remove(self.login_qr_code_path)
                    error("密匙错误")
                    return
                elif data == 86038:
                    os.remove(self.login_qr_code_path)
                    error("二维码已失效")
                    return
            time.sleep(2)

    def generate_qr_code(self, data: object = None) -> None:
        """
        生成二维码并打开
        :param data: 二维码数据
        :return:
        """
        if not data:
            error("data不能为空")
            return
        directory = os.path.dirname(self.login_qr_code_path)
        if not os.path.exists(directory):
            os.makedirs(directory)
        qr = QRCode(version=10, error_correction=constants.ERROR_CORRECT_H, box_size=15, border=4)
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image(image_factory=StyledPilImage, module_drawer=RoundedModuleDrawer(),
                            embeded_image_path=f"{os.getcwd()}\\icon\\logo.png")

        # 保存图片
        img.save(self.login_qr_code_path)
        # 打开图片
        Image.open(self.login_qr_code_path).show()

    @staticmethod
    def parser_url_params(url: str = "") -> dict:
        """
        解析URL参数,
        :param url:
        :return:
        """
        if not url:
            error("URL不能为哦空")
            return
        url_data = parse.urlparse(url)
        params = parse.parse_qs(url_data.query)
        result = {}
        for k, v in params.items():
            result[k] = v[0]
        return result

    def save_cookies(self, cookies: dict = {}) -> None:
        """
        cookies持久化
        :return:
        """
        cookies = cookies if cookies else self.cookies
        if not cookies:
            error("Cookies不能为空")
            return

        directory = os.path.dirname(self.cookies_save_path)
        if not os.path.exists(directory):
            os.makedirs(directory)

        with open(self.cookies_save_path, "wb") as file:
            pickle.dump(cookies, file)

    def read_cookies(self) -> dict:
        """
        从本地读取cookies
        :return:
        """
        if not os.path.exists(self.cookies_save_path):
            return
        with open(self.cookies_save_path, "rb") as file:
            cookies = pickle.load(file)

        self.cookies = cookies
        self.session.cookies.update(cookies)
        self.is_login = self.verify_cookies()
        if self.is_login:
            success("登陆成功")
        return cookies

    def verify_cookies(self, cookies: dict = {}):
        """
        验证cookies是否有效
        :return:
        """
        cookies = cookies if cookies else self.cookies
        if not cookies:
            error("Cookies不能为空")
            return
        self.session.cookies.update(cookies)
        result = self.get("api.bilibili.com/x/space/myinfo").json()
        if result['code'] == 0:
            self.is_vip = True if result['data']['vip']['status'] == 1 else False
            print(self.is_vip)
            print(120 if self.is_vip else 80 if self.is_login else 32)
            return True, result
        else:
            warn("Cookies验证失败, 当前未登录")
            os.remove(self.cookies_save_path)
            return False, None

    def vid(self, url: str):
        pattern = re.compile(r"/(\w+)/[?]")
        self.bvid = "".join(pattern.findall(url))
        return self.bvid

    def start_download(self):

        vid = self.vid(self.url)
        if not vid:
            error("获取视频ID错误")
            return

        result = self.get(f"api.bilibili.com/x/web-interface/view?bvid={vid}").json()
        if result['code'] != 0:
            error("获取视频信息错误")
            return

        data = result['data']
        # 获取视频详细信息
        self.title = data['title']
        # 验证路径
        self.save_path = self.valid_path()

        # 获取视频列表
        pages = data['pages']

        with self.progress:
            with ThreadPoolExecutor(max_workers=self.max_works) as pool:
                # 遍历视频列表
                for page in pages:
                    cid = str(page['cid'])
                    num = str(page['page'])
                    part = self.valid_title(str(page['part']))
                    part = part if part.startswith(num) else f"{num}. {part}"

                    # 获取视频信息
                    page_result = self.get(
                        f"api.bilibili.com/x/player/playurl?bvid={self.bvid}&cid={cid}&qn={127 if self.is_vip else 80 if self.is_login else 32}&fnval=0&fnver=0&fourk=1").json()
                    if page_result['code'] != 0:
                        error(f"{part}获取视频下载信息失败")
                        continue
                    # 视频信息
                    page_data = page_result['data']
                    # 视频保存路径
                    download_path = os.path.join(self.save_path, f"{part}.{page_data['format']}")
                    # 视频下载地址
                    download_url = page_data['durl'][0]['url']
                    # 下载标题
                    download_title = f"{part if len(part) < 10 else part[0:10] + '...'} {self.definition(page_data['quality'])}"
                    task_id = self.progress.add_task("download", filename=download_title, start=False)
                    pool.submit(self.download, download_url, download_path, task_id)
                    time.sleep(3)

    def download(self, url: str, path: str, task_id: TaskID):
        """
        下载视频
        :param task_id:
        :param url:
        :param path:
        :return:
        """

        response = urlopen(Request(url, headers=self.headers), timeout=self.timeout)
        self.progress.update(task_id, total=int(response.info()["Content-length"]))
        with open(path, "wb") as file:
            self.progress.stop_task(task_id)
            for data in iter(partial(response.read, 32768), b""):
                file.write(data)
                self.progress.update(task_id, advance=len(data))
                if self.done_event.is_set():
                    return

    def definition(self, num: int) -> str:
        """
        获取视频清晰度
        :param num:
        :return:
        """
        return self.definition_data[num]

    def valid_path(self, title: str = "") -> str:
        """
        验证路径是否合法
        :param title:
        :return:
        """
        title = title if title else self.title
        pattern = re.compile(r"[\\/:*?\"<>|]+")
        title = pattern.sub("", title)
        path = os.path.join(self.save_path, title)
        if not os.path.exists(path):
            os.makedirs(path)
        return path

    @staticmethod
    def valid_title(title: str) -> str:
        pattern = re.compile(r"[\\/:*?\"<>|]+")
        return pattern.sub("", title)

    def parse_cmd(self):
        # 从命令行获取参数
        parser = argparse.ArgumentParser(prefix_chars="-", description=self.name, epilog="帮助信息", add_help=False)
        parser.add_argument("-h", "--help", action="help", help="使用帮助")
        parser.add_argument("-u", "--url", dest="url", type=str, help="视频连接(必填)", required=False)
        parser.add_argument("-m", "--max", dest="max", type=int, help="同时下载数(默认: 5)", required=False, default=5)
        parser.add_argument("-e", "--exit", dest="exit", type=str, choices=["是", "否", "暂定"], help="完成是否关机(默认: 暂定)", required=False, default="暂定")
        parser.add_argument("-l", "--login", dest="login", type=str, choices=["是", "否", "暂定"], help="是否登录B站账户(默认: 是)", required=False, default="是")
        self.args = parser.parse_args()

    def exit(self):
        """
        是否关机
        :return:
        """
        if self.is_exit == "是":
            os.system("shutdown -s -t 10")
            return

    def __del__(self):
        """当实例被销毁时,释放掉session所持有的连接
        :return:
        """
        if self.session:
            self.session.close()


if __name__ == '__main__':
    BiliBiLi().login()